"""
    入口主程序
"""

import imu
from imu import deg2rad, rad2deg
import numpy as np
from math import sin, cos, tan
import complimentary_imu as cf
import matplotlib.pyplot as plt
from matplotlib import gridspec

dt = 0.01  # imu循环周期
global_x, global_y, global_vx, global_vy, vx_last, vy_last = 0, 0, 0, 0, 0, 0
local_vx, local_vy, local_vz = 0, 0, 0


def from_imu_get_position(attitude_hat: list, acc_data: list):
    """
    用IMU估计车辆的[x, y, vx, vy]
    :param attitude_hat:IMU融合后对姿态的估计值 [roll_hat, pitch_hat, yaw_hat] units:rad
    :param acc_data:加速度计的原始数据[ax, ay, az] units:m/s^2
    :return: [vx, vy, x, y] 返回纵向速度和横向速度（基于当前车体坐标），返回大地坐标系下的x, y值
    """
    global global_x, global_y, global_vx, global_vy, vx_last, vy_last
    global local_vx, local_vy, local_vz
    r = attitude_hat[0]
    p = attitude_hat[1]
    y = attitude_hat[2]
    # 将重力加速度从大地坐标系n系旋转到载体坐标系b系
    ax_temp = -sin(p) * 9.8
    ay_temp = sin(r) * cos(p) * 9.8
    az_temp = cos(r) * cos(p) * 9.8
    # 载体坐标系下的加速度减去重力加速度
    axb = acc_data[0] - ax_temp
    ayb = acc_data[1] - ay_temp
    azb = acc_data[2] - az_temp

    # 载体坐标系下的速度
    local_vx += dt * axb
    local_vy += dt * ayb
    local_vz += dt * azb
    # 全局坐标下的速度
    global_vx = local_vz * (sin(r) * sin(y) + cos(r) * cos(y) * sin(p)) - local_vy * (
                cos(r) * sin(y) - cos(y) * sin(p) * sin(r)) + local_vx * cos(p) * cos(y)
    global_vy = local_vy * (cos(r) * cos(y) + sin(p) * sin(r) * sin(y)) - local_vz * (
                cos(y) * sin(r) - cos(r) * sin(p) * sin(y)) + local_vx * cos(p) * sin(y)

    # 对全局速度进行一阶Runge-Kutta积分，得到全局位移
    global_x += dt * (global_vx + (global_vx - vx_last) / 2)
    global_y += dt * (global_vy + (global_vy - vy_last) / 2)
    vx_last = global_vx
    vy_last = global_vy
    return [global_x, global_y, local_vx, local_vy]


if __name__ == '__main__':
    imu = imu.IMU()
    complimentary_filter = cf.ComplimentaryFilter()
    imu.read_csv_file("imu_data1.csv")
    acc_hat_set = imu.get_acc_attitude_estimate()  # 计算加速度计的估计值[roll_0, pitch_0, roll_1, pitch_1, ...]
    gyro_data_set = imu.get_gyro_data()  # 获取陀螺仪去除零偏后的输出[gx_0', gy_0',gz_0','gx_1',.....]
    attitude_hat_set = []  # IMU姿态估计的结果
    position_hat_set = []  # 根据IMU姿态估计的结果，通过积分得到车辆的[x, y ,vx, vy]
    # 使用互补滤波算法，融合加速度计和陀螺仪的估计值
    for i in range(int(len(acc_hat_set) / 2)):
        acc_hat = [acc_hat_set[2 * i], acc_hat_set[2 * i + 1]]
        gyro_data = [gyro_data_set[3 * i], gyro_data_set[3 * i + 1], gyro_data_set[3 * i + 2]]
        # 互补滤波
        attitude_hat = complimentary_filter.complimentary_filter(acc_hat, gyro_data)
        acc_data = [imu.acc_data[3 * i], imu.acc_data[3 * i + 1], imu.acc_data[3 * i + 2]]
        position_hat_set += from_imu_get_position(attitude_hat, acc_data)  # [x, y, vx, vy]
        attitude_hat_set += attitude_hat

    # 画图数据准备
    # roll
    roll_plot = attitude_hat_set[0::3]
    roll_plot = [rad2deg(i) for i in roll_plot]
    roll_raw_plot = [rad2deg(i) for i in imu.roll_raw]

    # pitch
    pitch_plot = attitude_hat_set[1::3]
    pitch_plot = [rad2deg(i) for i in pitch_plot]
    pitch_raw_plot = [rad2deg(i) for i in imu.pitch_raw]

    # yaw
    yaw_plot = attitude_hat_set[2::3]
    yaw_plot = [rad2deg(i) for i in yaw_plot]
    yaw_raw_plot = [rad2deg(i) for i in imu.yaw_raw]

    # position
    x_plot = position_hat_set[0::4]
    y_plot = position_hat_set[1::4]
    vx_plot = position_hat_set[2::4]
    vy_plot = position_hat_set[3::4]

    # 姿态角画图
    plt.figure(figsize=(10, 6), dpi=160)
    plt.suptitle("IMU Raw Data VS 6-D IMU Complimentary Filter")
    gs = gridspec.GridSpec(7, 7)
    # roll angle
    ax11 = plt.subplot(gs[0:3, 0:3])
    plt.xlabel("time[s]")
    plt.ylabel("roll_raw[deg]")
    plt.grid()
    plt.plot(imu.t_raw_data, roll_raw_plot)

    ax12 = plt.subplot(gs[0:3, 4:])
    plt.xlabel("time[s]")
    plt.ylabel("roll[deg]")
    plt.grid()
    plt.plot(imu.t_raw_data, roll_plot)

    # pitch angle
    ax13 = plt.subplot(gs[4:, 0:3])
    plt.xlabel("time[s]")
    plt.ylabel("pitch_raw[deg]")
    plt.grid()
    plt.plot(imu.t_raw_data, pitch_raw_plot)

    ax14 = plt.subplot(gs[4:, 4:])
    plt.xlabel("time[s]")
    plt.ylabel("pitch[deg]")
    plt.grid()
    plt.plot(imu.t_raw_data, pitch_plot)

    # yaw angle
    plt.figure(figsize=(10, 6), dpi=160)
    plt.suptitle("IMU Raw Data VS 6-D IMU Complimentary Filter")
    gs = gridspec.GridSpec(1, 7)

    ax21 = plt.subplot(gs[0, 0:3])
    plt.xlabel("time[s]")
    plt.ylabel("yaw_raw[deg]")
    plt.grid()
    plt.plot(imu.t_raw_data, yaw_raw_plot)

    ax22 = plt.subplot(gs[0, 4:])
    plt.xlabel("time[s]")
    plt.ylabel("yaw[deg]")
    plt.grid()
    plt.plot(imu.t_raw_data, yaw_plot)

    # 画轨迹
    plt.figure(figsize=(10, 6), dpi=160)
    plt.suptitle("Position Plot")
    gs = gridspec.GridSpec(7, 7)
    # vx-t
    ax31 = plt.subplot(gs[0:3, 0:3])
    plt.xlabel("time[s]")
    plt.ylabel("vx[m/s]")
    plt.grid()
    plt.plot(imu.t_raw_data, vx_plot)
    # vy-t
    ax32 = plt.subplot(gs[0:3, 4:])
    plt.xlabel("time[s]")
    plt.ylabel("vy[m/s]")
    plt.grid()
    plt.plot(imu.t_raw_data, vy_plot)

    # x-y
    ax33 = plt.subplot(gs[4:, 2:5])
    plt.xlabel("x[m]")
    plt.ylabel("y[m]")
    plt.grid()
    plt.plot(x_plot, y_plot)

    plt.show()
